class My_VGG16(nn.Module):
    def __init__(self,class_num):
        super(My_VGG16, self).__init__()
        # 特征提取层
        self.features = nn.Sequential(
            nn.Conv2d(in_channels=3,
                       out_channels=64,
                       kernel_size=3,
                       stride=1,
                       padding=1),
            nn.Conv2d(64, 64, 3, 1, 1),
            nn.MaxPool2d(kernel_size=2,stride=2),
            nn.Conv2d(64, 128, 3, 1, 1),
            nn.Conv2d(128, 128, 3, 1, 1),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(128, 256, 3, 1, 1),
            nn.Conv2d(256, 256, 3, 1, 1),
            nn.Conv2d(256, 256, 3, 1, 1),
            
            nn.MaxPool2d(kernel_size=2,stride=2),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.Conv2d(512, 512, 3, 1, 1),
            nn.Conv2d(512, 512, 3, 1, 1),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(512, 512, 3, 1, 1),
            nn.Conv2d(512, 512, 3, 1, 1),
            nn.Conv2d(512, 512, 3, 1, 1),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        # 分类层
        self.classifier = nn.Sequential(
                               nn.Linear(in_features=1*1*512, out_features=4096),
                               nn.ReLU(),
                               nn.Dropout(0.5),
                              nn.Linear(4096,4096),
                              nn.ReLU(),
                              nn.Dropout(0.5),
                              nn.Linear(4096,class_num)
                             )

    def forward(self,x):
        x = self.features(x)
        x = torch.flatten(x,1)
        result = self.classifier(x)
        result = nn.Softmax(dim=-1)(result)
        return result 
